import os
import glob
import argparse
from pod5 import Reader as Pod5Reader
from pod5 import  Writer as Pod5Writer
import pod5
import numpy as np

def main():
    parser = argparse.ArgumentParser(description="Extract reads from multiple pod5 files based on a read ID list and split them into chunks.")
    parser.add_argument("-i", "--input_dir", required=True, help="Input directory containing .pod5 files.")
    # parser.add_argument("-o", "--output_dir", required=True, help="Output directory to store extracted pod5 files.")
    parser.add_argument("-r", "--read_id_file",default=None, required=False, help="File containing read IDs to extract, one per line.")
    
    args = parser.parse_args()

    input_dir = args.input_dir
    # output_dir = args.output_dir
    read_id_file = args.read_id_file
    

    # # 创建输出目录（如果不存在）
    # if not os.path.exists(output_dir):
    #     os.makedirs(output_dir)

    # 读取需要提取的read_id列表
    read_ids_to_extract=None
    if read_id_file is not None:
        read_ids_to_extract = set(line.strip() for line in open(read_id_file, "r") if line.strip())
        print(f"Read IDs to extract: {len(read_ids_to_extract)}")

    # 找到输入目录中所有的pod5文件
    pod5_files = glob.glob(os.path.join(input_dir, "**", "*.pod5"), recursive=True)
    if not pod5_files:
        print("No .pod5 files found in the input directory.")
        return

    current_batch = []
    batch_num = 1
    run_info_copied = False
    run_info_example = None
    process_num=0
    np.set_printoptions(threshold=np.inf)
    # 逐个pod5文件处理
    for pod5_file in pod5_files:
        with Pod5Reader(pod5_file) as reader:
            # 保存run_info以用于后续写出（确保多个pod5文件兼容）
            if not run_info_copied:
                
                run_info_copied = True

            for read in reader.reads():
                if read_ids_to_extract is not None and str(read.read_id) not in read_ids_to_extract:
                    continue
                process_num+=1
                signal = read.signal
                read_id=read.read_id
                print(f"Processing read {read_id}...")
                print(f"Signal: {np.array(signal)}")

    # 处理最后不满chunk_size的剩余reads
    print(f"Processed {process_num} reads.")

    print("Extraction and splitting completed successfully.")

if __name__ == "__main__":
    main()
